import asyncio
from typing import Optional, Dict, Type, Tuple

from ofrak.model.tag_model import ResourceTag

from ofrak.component.unpacker import Unpacker
from ofrak.core.code_region import CodeRegion
from ofrak.model.resource_model import ResourceAttributes
from ofrak.model.viewable_tag_model import AttributesType
from ofrak.resource import Resource
from ofrak.service.resource_service_i import ResourceFilter
from ofrak.core.elf.model import (
    ElfHeader,
    ElfProgramHeader,
    ElfProgramHeaderType,
    ElfSegment,
    ElfSectionHeader,
    ElfSection,
    ElfStringSection,
    ElfSectionType,
    ElfSymbolSection,
    ElfSymbol,
    ElfBasicHeader,
    ElfSectionFlag,
    ElfSectionNameStringSection,
    Elf,
    ElfSegmentStructure,
    ElfSectionStructure,
    ElfRelaSection,
    ElfRelaEntry,
    ElfDynSymbolSection,
    ElfDynamicSection,
    ElfDynamicEntry,
    ElfFiniArraySection,
    ElfInitArraySection,
    ElfPointerArraySection,
    ElfSymbolStructure,
    ElfVirtualAddress,
)
from ofrak_type.bit_width import BitWidth
from ofrak_type.memory_permissions import MemoryPermissions
from ofrak_type.range import Range


class ElfUnpacker(Unpacker[None]):
    """
    Extracts and parses ELF (Executable and Linkable Format) binary structures including the basic header, full header with entry point and architecture information, program headers describing loadable segments, the segments themselves (like PT_LOAD, PT_DYNAMIC), section headers, sections (.text, .data, .bss, etc.), string sections, symbol tables, and executable code regions. Use when analyzing any Linux/Unix executable, shared library, object file, or kernel module to examine code, data, symbols, and linking information.
    """

    id = b"ElfUnpacker"
    targets = (Elf,)
    children = (
        ElfBasicHeader,
        ElfHeader,
        ElfProgramHeader,
        ElfSegment,
        ElfSectionHeader,
        ElfSection,
        ElfStringSection,
        ElfSymbolSection,
        CodeRegion,
    )

    async def unpack(self, resource: Resource, config=None):
        """
        Unpack ELF headers and sections / segments into the OFRAK resource tree.

        After unpacking the ElfHeader, the unpacker first unpacks ElfSections defined in the ElfSectionHeader into the
        resource tree.

        If no executable sections are found, then the proceeding ElfProgramHeader unpacking routine is permitted to
        unpack LOAD-able ElfSegments into the resource tree. Executable LOAD-able ElfSegments are unpacked into
        CodeRegions.

        legend:  +    composable child
                 ^    exclusive child
                [ ]   optional / conditional child
                (s)   one or more instances of child

        RESOURCE
         + ElfBasicHeader
         + ElfHeader
           [+] ElfProgramHeader              (creates LOAD-able ElfSegments if ElfSectionHeaders yields no CodeRegions)
           [+] ElfSectionHeader
        [+] ElfSection(s)
           [^] ElfFiniArraySection           (if FINI_ARRAY)
           [^] ElfInitArraySection           (if INIT_ARRAY)
           [^] ElfDynamicSection             (if DYNAMIC)
           [^] ElfRelaSection                (if RELA)
           [^] ElfDynSymbolSection           (if DYNSYM)
           [^] ElfSymbolSection              (if SYMTAB)
           [^] ElfStringSection              (if STRTAB)
           [+] CodeRegion                    (if executable)
           [+] ElfSectionNameStringSection   (if name string section)
        [+] ElfSegment(s)                    (if ElfSection yields no CodeRegions,
           [+] CodeRegion (if executable)     LOAD-able ElfSegments will populate the body of the resource tree instead)
        """

        e_basic_header_r = await resource.create_child(
            tags=(ElfBasicHeader,), data_range=Range(0, 16)
        )
        e_basic_header = await e_basic_header_r.view_as(ElfBasicHeader)

        e_header_range = Range.from_size(
            16, 36 if e_basic_header.get_bitwidth() is BitWidth.BIT_32 else 48
        )
        e_header_r = await resource.create_child(tags=(ElfHeader,), data_range=e_header_range)
        e_header = await e_header_r.view_as(ElfHeader)

        ###########################################################################################
        ### Unpack section headers and associated sections

        sections_by_range_start: Dict[int, Resource] = dict()
        code_region_present = None

        # Create the section header/body resources
        for index in range(e_header.e_shnum):
            e_section_header_offset = e_header.e_shoff + index * e_header.e_shentsize
            e_section_header_range = Range.from_size(e_section_header_offset, e_header.e_shentsize)
            e_section_header_r = await resource.create_child(
                tags=(ElfSectionHeader,),
                data_range=e_section_header_range,
                attributes=(AttributesType[ElfSectionStructure](index),),
            )
            e_section_header = await e_section_header_r.view_as(ElfSectionHeader)

            e_section_offset = e_section_header.sh_offset
            e_section_range = Range.from_size(e_section_offset, e_section_header.sh_size)
            opt_e_section_range: Optional[Range] = e_section_range
            if e_section_range.length() == 0:
                # It's possible that the section does not contain any data
                opt_e_section_range = None
            if e_section_header.get_type() is ElfSectionType.NOBITS:
                # NOBITS sections never have data in the file; their sh_size refers to in-mem size
                opt_e_section_range = None
            e_section_r = await resource.create_child(
                tags=(ElfSection,),
                data_range=opt_e_section_range,
                attributes=(AttributesType[ElfSectionStructure](index),),
            )
            sections_by_range_start[e_section_offset] = e_section_r
            if e_section_header.get_type() is ElfSectionType.FINI_ARRAY:
                e_section_r.add_tag(ElfFiniArraySection)
            elif e_section_header.get_type() is ElfSectionType.INIT_ARRAY:
                e_section_r.add_tag(ElfInitArraySection)
            elif e_section_header.get_type() is ElfSectionType.DYNAMIC:
                e_section_r.add_tag(ElfDynamicSection)
            elif e_section_header.get_type() is ElfSectionType.RELA:
                e_section_r.add_tag(ElfRelaSection)
            elif e_section_header.get_type() is ElfSectionType.DYNSYM:
                e_section_r.add_tag(ElfDynSymbolSection)
            elif e_section_header.get_type() is ElfSectionType.SYMTAB:
                e_section_r.add_tag(ElfSymbolSection)
            elif e_section_header.get_type() is ElfSectionType.STRTAB:
                e_section_r.add_tag(ElfStringSection)

            if e_section_header.has_flag(ElfSectionFlag.EXECINSTR):
                e_section_r.add_tag(CodeRegion)
                code_region_present = True

            if index == e_header.e_shstrndx:
                e_section_r.add_tag(ElfSectionNameStringSection)

        ###########################################################################################
        ### Unpack segment headers, conditionally unpack associated loadable segments

        # Create the program header resources
        for index in range(e_header.e_phnum):
            e_program_header_offset = e_header.e_phoff + (index * e_header.e_phentsize)
            e_program_header_range = Range.from_size(e_program_header_offset, e_header.e_phentsize)

            e_program_header_r = await resource.create_child(
                tags=(ElfProgramHeader,),
                data_range=e_program_header_range,
                attributes=(AttributesType[ElfSegmentStructure](index),),
            )

            e_program_header = await e_program_header_r.view_as(ElfProgramHeader)

            # Don't unpack loadable segments if the section unpacker had already unpacked any CodeRegions
            if code_region_present:
                continue

            if e_program_header.p_type == ElfProgramHeaderType.LOAD.value:
                e_segment_offset = e_program_header.p_offset
                e_segment_range = Range.from_size(e_segment_offset, e_program_header.p_filesz)
                opt_e_segment_range: Optional[Range] = e_segment_range

                # Loaded segments can have no data to initialize with (heap, bss, etc.)
                if e_segment_range.length() == 0:
                    opt_e_segment_range = None

                # We need to inform OFRAK data service the order of child nodes to be populated into the resource tree,
                # since there may be overlapping regions due to flattening non-flat memory structures derived from the
                # binary-under-analysis. We should deprecate `data_after` and `data_before` in favor of a structure
                # that can hold multiple memory views (virtual, physical, file, etc.) of an analyzed binary.
                e_segment_r = await resource.create_child(
                    tags=(ElfSegment,),
                    data_range=opt_e_segment_range,
                    attributes=(AttributesType[ElfSegmentStructure](index),),
                )

                # Tag the segment as a CodeRegion if the loaded segment is executable
                memory_permissions = e_program_header.get_memory_permissions()
                if memory_permissions & MemoryPermissions.R is MemoryPermissions.R:
                    e_segment_r.add_tag(CodeRegion)
                await e_segment_r.save()


class ElfDynamicSectionUnpacker(Unpacker[None]):
    """
    Extracts individual dynamic linking entries from an ELF binary's .dynamic section, which
    controls runtime linking behavior. Each entry is a tag-value pair that specifies things like
    required shared libraries, symbol/string table locations, relocation table addresses,
    initialization/finalization functions, and runtime linker options. Use when you need to
    understand an ELF's runtime dependencies, analyze how dynamic linking works, or prepare to
    modify dynamic linking behavior. Essential for analyzing shared libraries, dynamically linked
    executables, and understanding the runtime loader's configuration.
    """

    id = b"ElfDynamicSectionUnpacker"
    targets = (ElfDynamicSection,)
    children = (ElfDynamicEntry,)

    async def unpack(self, resource: Resource, config=None):
        e_section = await resource.view_as(ElfDynamicSection)
        elf_r = await e_section.get_elf()
        e_basic_header = await elf_r.get_basic_header()
        dyn_entry_size = 16 if e_basic_header.get_bitwidth() is BitWidth.BIT_64 else 8
        await make_children_helper(resource, ElfDynamicEntry, dyn_entry_size, None)


class ElfRelaUnpacker(Unpacker[None]):
    """
    Extracts individual relocation entries with addends from ELF relocation sections (.rela.text,
    .rela.dyn, etc.). Relocation entries specify how addresses should be adjusted when code is
    loaded at different addresses or linked with other objects. Use when analyzing
    position-independent code (PIC), understanding dynamic linking relocations, or preparing to
    modify code that requires relocation adjustments. Each entry contains an offset, relocation
    type, symbol reference, and addend value.
    """

    id = b"ElfRelaUnpacker"
    targets = (ElfRelaSection,)
    children = (ElfRelaEntry,)

    async def unpack(self, resource: Resource, config=None):
        e_section = await resource.view_as(ElfRelaSection)
        elf_r = await e_section.get_elf()
        e_basic_header = await elf_r.get_basic_header()
        rela_size = 24 if e_basic_header.get_bitwidth() is BitWidth.BIT_64 else 12
        await make_children_helper(resource, ElfRelaEntry, rela_size, None)


class ElfSymbolUnpacker(Unpacker[None]):
    """
    Extracts individual symbol entries from ELF symbol tables (.symtab or .dynsym), where each
    symbol represents a function, variable, or other named entity with an address, size, and type.
    These symbols reflect what the ELF header claims, which might not match what's actually in the
    binary (e.g., after code modifications, or in stripped/obfuscated binaries). Use when you need
    to modify the symbol table that the OS loader will read, but don't rely on these for discovering
    actual symbols - use analysis-based symbol discovery instead.
    """

    id = b"ElfSymbolUnpacker"
    targets = (ElfSymbolSection,)
    children = (ElfSymbol,)

    async def unpack(self, resource: Resource, config=None):
        e_section = await resource.view_as(ElfSymbolSection)
        elf_r = await e_section.get_elf()
        e_basic_header = await elf_r.get_basic_header()
        symbol_size = 16 if e_basic_header.get_bitwidth() is BitWidth.BIT_32 else 24
        await make_children_helper(
            resource, ElfSymbol, symbol_size, AttributesType[ElfSymbolStructure]
        )


class ElfPointerArraySectionUnpacker(Unpacker[None]):
    """
    Extracts individual pointer values from ELF sections that contain arrays of function pointers,
    such as .init_array (constructors run before main), .fini_array (destructors run after main),
    .ctors, and .dtors. Each pointer references a function that should be called during
    initialization or cleanup. Use when analyzing constructor/destructor functions, understanding
    initialization order, or modifying the functions called during program startup/shutdown.
    """

    id = b"ElfPointerArraySectionUnpacker"
    targets = (ElfPointerArraySection,)
    children = (ElfVirtualAddress,)

    async def unpack(self, resource: Resource, config=None):
        elf_r = await resource.get_only_ancestor_as_view(Elf, ResourceFilter.with_tags(Elf))
        e_basic_header = await elf_r.get_basic_header()
        addr_size = 4 if e_basic_header.get_bitwidth() is BitWidth.BIT_32 else 8
        await make_children_helper(resource, ElfVirtualAddress, addr_size, None)


async def make_children_helper(
    resource: Resource,
    entry_type: ResourceTag,
    entry_size: int,
    structure_index_type: Optional[Type[ResourceAttributes]],
) -> None:
    elf_section_size = await resource.get_data_length()
    create_child_tasks = []
    for i, offset in enumerate(range(0, elf_section_size, entry_size)):
        if structure_index_type is not None:
            attrs: Tuple[ResourceAttributes, ...] = (structure_index_type(i),)  # type: ignore
        else:
            attrs = ()
        create_child_tasks.append(
            resource.create_child(
                tags=(entry_type,),
                data_range=Range.from_size(offset, entry_size),
                attributes=attrs,
            )
        )
    await asyncio.gather(*create_child_tasks)
